"""
MNE数据处理器

本文件用于基于MNE库的脑波数据预处理和滤波功能，包括：
- 带通滤波器：去除低频漂移和高频噪声
- 陷波滤波器：抑制工频噪声(50Hz/60Hz)
- 数据预处理：格式转换和质量检查
- 滤波器参数优化：自适应参数调整

使用方法：
- 使用MNE库的内置滤波算法，确保信号处理的专业性和可靠性
- 支持多种滤波器设计方法(FIR/IIR)
- 提供滤波前后的信号质量评估

原理：
通过MNE库的专业滤波算法，有效去除脑波信号中的各种噪声和伪迹，
为后续的YASA睡眠分析提供高质量的预处理数据。
"""

import numpy as np
from typing import Optional, Union, Tuple, Dict, List
import warnings
from datetime import datetime
from scipy import signal

# 尝试导入MNE库，如果不存在则使用模拟实现
try:
    import mne
    MNE_AVAILABLE = True
except ImportError:
    MNE_AVAILABLE = False
    # 创建MNE模拟类
    class MockMNE:
        @staticmethod
        def set_log_level(level):
            pass
        
        @staticmethod
        def create_info(ch_names, sfreq, ch_types, verbose=False):
            return {'ch_names': ch_names, 'sfreq': sfreq, 'ch_types': ch_types}
        
        class io:
            @staticmethod
            class RawArray:
                def __init__(self, data, info, verbose=False):
                    self.data = data
                    self.info = info
                    self._sfreq = info['sfreq']
                
                def copy(self):
                    return MockMNE.io.RawArray(self.data.copy(), self.info)
                
                def get_data(self):
                    return self.data
                
                def filter(self, l_freq=None, h_freq=None, fir_design='firwin', 
                          filter_length='auto', phase='zero', verbose=False):
                    # 使用scipy实现简单的带通滤波
                    if l_freq is not None and h_freq is not None:
                        nyquist = self._sfreq / 2
                        low = l_freq / nyquist
                        high = h_freq / nyquist
                        b, a = signal.butter(4, [low, high], btype='band')
                        for i in range(self.data.shape[0]):
                            self.data[i] = signal.filtfilt(b, a, self.data[i])
                
                def notch_filter(self, freqs, fir_design='firwin', 
                               filter_length='auto', phase='zero', verbose=False):
                    # 使用scipy实现简单的陷波滤波
                    if isinstance(freqs, (int, float)):
                        freqs = [freqs]
                    
                    for freq in freqs:
                        nyquist = self._sfreq / 2
                        # 创建陷波滤波器
                        Q = 30  # 品质因子
                        w0 = freq / nyquist
                        b, a = signal.iirnotch(w0, Q)
                        for i in range(self.data.shape[0]):
                            self.data[i] = signal.filtfilt(b, a, self.data[i])
        
        __version__ = "0.0.0-mock"
    
    mne = MockMNE()

from app.models.brain_wave_data import BrainWaveData, TimeSeriesData, FrequencyData
from app.exceptions.device_exceptions import DeviceDataException
import logging

# 配置日志
logger = logging.getLogger(__name__)


class MNEProcessor:
    """MNE数据处理器类"""
    
    def __init__(self, sampling_rate: float = 512.0, verbose: bool = False):
        """
        初始化MNE处理器
        
        Args:
            sampling_rate: 采样率 (Hz)，默认512Hz适配NeuroSky设备
            verbose: 是否显示详细日志
        """
        self.sampling_rate = sampling_rate
        self.verbose = verbose
        self.mne_available = MNE_AVAILABLE
        
        # 设置MNE日志级别
        if not verbose:
            mne.set_log_level('WARNING')
        
        # 针对512Hz采样率和清醒/睡眠区分优化的滤波参数
        self.default_bandpass_params = {
            'l_freq': 2.0,      # 低频截止频率 - 更好地去除低频噪声和基线漂移，提高清醒/睡眠区分准确性
            'h_freq': 40.0,     # 高频截止频率 - 去除肌电干扰，保留脑波
            'fir_design': 'firwin',  # FIR滤波器设计方法
            'filter_length': 'auto',  # 让MNE自动计算，但会在_adjust_filter_params中调整
            'phase': 'zero',    # 零相位滤波
            'l_trans_bandwidth': 2.0,  # 低频过渡带宽度，设置为固定值避免自动计算过小
            'h_trans_bandwidth': 10.0   # 高频过渡带宽度，设置为固定值
        }
        
        self.default_notch_params = {
            'freqs': [50.0, 60.0],  # 陷波频率(工频噪声)
            'fir_design': 'firwin',
            'filter_length': 'auto',  # 让MNE自动计算，但会在_adjust_filter_params中调整
            'phase': 'zero',
            'trans_bandwidth': 'auto'  # 让MNE自动计算，但会在_adjust_filter_params中调整
        }
    
    def create_raw_from_data(self, data: np.ndarray, ch_names: Optional[List[str]] = None) -> mne.io.RawArray:
        """
        从数据数组创建MNE Raw对象
        
        Args:
            data: 脑波数据数组，形状为 (n_channels, n_samples)
            ch_names: 通道名称列表
            
        Returns:
            MNE Raw对象
        """
        # 确保数据是二维数组
        if data.ndim == 1:
            data = data.reshape(1, -1)
        
        n_channels = data.shape[0]
        
        # 设置默认通道名称
        if ch_names is None:
            ch_names = [f'EEG_{i+1}' for i in range(n_channels)]
        
        # 创建Info对象
        info = mne.create_info(
            ch_names=ch_names,
            sfreq=self.sampling_rate,
            ch_types=['eeg'] * n_channels,
            verbose=self.verbose
        )
        
        # 创建Raw对象
        raw = mne.io.RawArray(data, info, verbose=self.verbose)
        
        return raw
    
    def apply_bandpass_filter(self, raw: mne.io.Raw, l_freq: float = None, h_freq: float = None, **kwargs) -> mne.io.Raw:
        """
        应用带通滤波器
        
        Args:
            raw: MNE Raw对象
            l_freq: 低频截止频率 (Hz)
            h_freq: 高频截止频率 (Hz)
            **kwargs: 其他滤波参数
            
        Returns:
            滤波后的Raw对象
        """
        # 使用默认参数或用户指定参数
        params = self.default_bandpass_params.copy()
        if l_freq is not None:
            params['l_freq'] = l_freq
        if h_freq is not None:
            params['h_freq'] = h_freq
        params.update(kwargs)
        
        # 获取信号长度并调整滤波器参数
        n_samples = raw.n_times
        params = self._adjust_filter_params(params, n_samples)
        
        # 创建副本以避免修改原始数据
        filtered_raw = raw.copy()
        
        # 应用带通滤波器 - 使用MNE内置滤波算法
        # l_freq: 低频截止点，去除低频漂移和直流分量
        # h_freq: 高频截止点，去除高频噪声和肌电伪迹
        # fir_design: 使用firwin方法设计FIR滤波器，提供线性相位响应
        # filter_length: 动态调整滤波器长度，确保不超过信号长度
        # phase: 零相位滤波，避免信号时间延迟
        filtered_raw.filter(
            l_freq=params['l_freq'],
            h_freq=params['h_freq'],
            fir_design=params['fir_design'],
            filter_length=params['filter_length'],
            phase=params['phase'],
            verbose=self.verbose
        )
        
        return filtered_raw
    
    def apply_notch_filter(self, raw: mne.io.Raw, freqs: Union[float, List[float]] = None, **kwargs) -> mne.io.Raw:
        """
        应用陷波滤波器
        
        Args:
            raw: MNE Raw对象
            freqs: 陷波频率或频率列表 (Hz)
            **kwargs: 其他滤波参数
            
        Returns:
            滤波后的MNE Raw对象
        """
        # 使用默认参数，允许用户覆盖
        params = self.default_notch_params.copy()
        if freqs is not None:
            params['freqs'] = freqs if isinstance(freqs, list) else [freqs]
        params.update(kwargs)
        
        # 获取信号长度并调整滤波器参数
        n_samples = raw.n_times
        params = self._adjust_filter_params(params, n_samples)
        
        # 创建副本以避免修改原始数据
        filtered_raw = raw.copy()
        
        # 应用陷波滤波器 - 使用MNE内置陷波滤波算法
        # freqs: 需要抑制的频率列表，通常是50Hz或60Hz工频噪声
        # fir_design: 使用firwin方法设计陷波滤波器
        # filter_length: 动态调整滤波器长度，确保不超过信号长度
        # phase: 零相位滤波，保持信号的时间特性
        # trans_bandwidth: 过渡带宽度，影响滤波器的陡峭程度
        filtered_raw.notch_filter(
            freqs=params['freqs'],
            fir_design=params['fir_design'],
            filter_length=params['filter_length'],
            phase=params['phase'],
            trans_bandwidth=params['trans_bandwidth'],
            verbose=self.verbose
        )
        
        return filtered_raw
    
    def _adjust_filter_params(self, params: Dict, n_samples: int) -> Dict:
        """
        根据信号长度调整滤波器参数，针对512Hz采样率优化
        
        Args:
            params: 滤波器参数字典
            n_samples: 信号样本数
            
        Returns:
            调整后的参数字典
        """
        adjusted_params = params.copy()
        
        # 根据MNE文档，滤波器长度必须足够满足过渡带要求
        # 对于很短的信号，我们需要大幅放宽过渡带参数
        if n_samples < 512:  # 对于少于1秒的极短信号
            # 极短信号：使用最宽松的过渡带参数
            if 'l_trans_bandwidth' in adjusted_params:
                if adjusted_params['l_trans_bandwidth'] == 'auto' or isinstance(adjusted_params['l_trans_bandwidth'], str):
                    adjusted_params['l_trans_bandwidth'] = 10.0  # 10Hz过渡带，非常宽松
            if 'h_trans_bandwidth' in adjusted_params:
                if adjusted_params['h_trans_bandwidth'] == 'auto' or isinstance(adjusted_params['h_trans_bandwidth'], str):
                    adjusted_params['h_trans_bandwidth'] = 15.0  # 15Hz过渡带，非常宽松
            if 'trans_bandwidth' in adjusted_params:
                if adjusted_params['trans_bandwidth'] == 'auto' or isinstance(adjusted_params['trans_bandwidth'], str):
                    adjusted_params['trans_bandwidth'] = 10.0   # 10Hz过渡带，适合陷波滤波器
            
            # 对于极短信号，使用最小可能的滤波器长度
            if adjusted_params['filter_length'] == 'auto':
                adjusted_params['filter_length'] = max(32, n_samples // 8)  # 最小32，或信号长度的1/8
            elif isinstance(adjusted_params['filter_length'], int):
                adjusted_params['filter_length'] = min(adjusted_params['filter_length'], max(32, n_samples // 8))
                
        elif n_samples < 2048:  # 对于少于4秒的短信号
            # 短信号：使用宽松的过渡带参数
            if 'l_trans_bandwidth' in adjusted_params:
                if adjusted_params['l_trans_bandwidth'] == 'auto' or isinstance(adjusted_params['l_trans_bandwidth'], str):
                    adjusted_params['l_trans_bandwidth'] = 5.0  # 5Hz过渡带
            if 'h_trans_bandwidth' in adjusted_params:
                if adjusted_params['h_trans_bandwidth'] == 'auto' or isinstance(adjusted_params['h_trans_bandwidth'], str):
                    adjusted_params['h_trans_bandwidth'] = 10.0  # 10Hz过渡带
            if 'trans_bandwidth' in adjusted_params:
                if adjusted_params['trans_bandwidth'] == 'auto' or isinstance(adjusted_params['trans_bandwidth'], str):
                    adjusted_params['trans_bandwidth'] = 8.0   # 8Hz过渡带
            
            # 计算最大可用的滤波器长度
            max_filter_length = max(64, n_samples // 6)
            if adjusted_params['filter_length'] == 'auto':
                adjusted_params['filter_length'] = min(max_filter_length, 512)
            elif isinstance(adjusted_params['filter_length'], int):
                adjusted_params['filter_length'] = min(adjusted_params['filter_length'], max_filter_length)
                
        else:
            # 对于长信号（>=4秒），可以使用更严格的过渡带参数
            if 'l_trans_bandwidth' in adjusted_params:
                if adjusted_params['l_trans_bandwidth'] == 'auto' or isinstance(adjusted_params['l_trans_bandwidth'], str):
                    adjusted_params['l_trans_bandwidth'] = 1.0  # 1Hz过渡带，更严格
            if 'h_trans_bandwidth' in adjusted_params:
                if adjusted_params['h_trans_bandwidth'] == 'auto' or isinstance(adjusted_params['h_trans_bandwidth'], str):
                    adjusted_params['h_trans_bandwidth'] = 5.0  # 5Hz过渡带
            if 'trans_bandwidth' in adjusted_params:
                if adjusted_params['trans_bandwidth'] == 'auto' or isinstance(adjusted_params['trans_bandwidth'], str):
                    adjusted_params['trans_bandwidth'] = 2.0   # 2Hz过渡带，适合陷波滤波器
            
            # 允许更长的滤波器长度
            if adjusted_params['filter_length'] == 'auto':
                max_allowed_length = min(8192, n_samples // 3)  # 最大8192，或信号长度的1/3
                adjusted_params['filter_length'] = max_allowed_length
            elif isinstance(adjusted_params['filter_length'], int):
                adjusted_params['filter_length'] = min(adjusted_params['filter_length'], n_samples // 3)
        
        return adjusted_params
    
    def apply_combined_filter(self, raw: mne.io.Raw, 
                            l_freq: float = None, h_freq: float = None,
                            notch_freqs: Union[float, List[float]] = None,
                            **kwargs) -> mne.io.Raw:
        """
        应用组合滤波器（带通 + 陷波）
        
        Args:
            raw: MNE Raw对象
            l_freq: 带通滤波低频截止频率
            h_freq: 带通滤波高频截止频率
            notch_freqs: 陷波频率
            **kwargs: 其他滤波参数
            
        Returns:
            滤波后的Raw对象
        """
        # 先应用带通滤波器
        # 带通滤波器首先去除信号中的低频漂移和高频噪声
        # 这为后续的陷波滤波提供更干净的基础信号
        filtered_raw = self.apply_bandpass_filter(raw, l_freq, h_freq, **kwargs)
        
        # 再应用陷波滤波器
        # 陷波滤波器专门针对工频噪声进行精确抑制
        # 在带通滤波基础上进一步提高信号质量
        if notch_freqs is not None:
            filtered_raw = self.apply_notch_filter(filtered_raw, notch_freqs, **kwargs)
        
        return filtered_raw
    
    def process_brain_wave_data(self, brain_wave_data: BrainWaveData, 
                               filter_type: str = 'combined',
                               **filter_params) -> BrainWaveData:
        """
        处理BrainWaveData对象
        
        Args:
            brain_wave_data: 脑波数据对象
            filter_type: 滤波类型 ('bandpass', 'notch', 'combined')
            **filter_params: 滤波参数
            
        Returns:
            处理后的BrainWaveData对象
        """
        if not brain_wave_data.time_series:
            raise DeviceDataException("没有时间序列数据可供处理")
        
        # 提取时间序列数据
        time_data = brain_wave_data.time_series
        raw_values = time_data.data
        
        # 检查数据长度是否足够进行滤波分析
        min_samples = 256  # 最小样本数要求
        if len(raw_values) < min_samples:
            logger.warning(f"信号长度 ({len(raw_values)}) 小于推荐的最小长度 ({min_samples})，可能影响滤波效果")
            # 对于过短的信号，可以选择跳过滤波或使用更简单的处理方法
            if len(raw_values) < 64:
                logger.warning("信号过短，跳过滤波处理")
                return brain_wave_data
        
        # 创建MNE Raw对象
        raw = self.create_raw_from_data(raw_values)
        
        # 根据滤波类型应用相应的滤波器
        try:
            if filter_type == 'bandpass':
                # 仅应用带通滤波器
                filtered_raw = self.apply_bandpass_filter(raw, **filter_params)
            elif filter_type == 'notch':
                # 仅应用陷波滤波器
                filtered_raw = self.apply_notch_filter(raw, **filter_params)
            elif filter_type == 'combined':
                # 应用组合滤波器（推荐）
                filtered_raw = self.apply_combined_filter(raw, **filter_params)
            else:
                raise ValueError(f"不支持的滤波类型: {filter_type}")
        except Exception as e:
            logger.error(f"滤波处理失败: {e}")
            # 如果滤波失败，返回原始数据
            return brain_wave_data
        
        # 获取滤波后的数据
        filtered_data = filtered_raw.get_data()[0]  # 假设单通道
        
        # 创建新的时间序列数据对象，保持时间戳信息
        filtered_time_data = TimeSeriesData(
            data=filtered_data,
            sampling_rate=time_data.sampling_rate,
            timestamp=time_data.timestamp,  # 保持原始时间戳
            channel_names=time_data.channel_names
        )
        
        # 创建新的BrainWaveData对象，添加处理时间戳
        filtered_brain_wave_data = BrainWaveData(
            time_series=filtered_time_data,
            frequency_data=brain_wave_data.frequency_data,  # 保持原有频域数据
            session_id=brain_wave_data.session_id,
            processing_stage="filtered",  # 修复：统一使用"filtered"标记
            processing_timestamp=datetime.now()  # 添加处理时间戳
        )
        
        return filtered_brain_wave_data
    
    async def process_single_sample(self, brain_wave_data: BrainWaveData) -> BrainWaveData:
        """
        处理单个脑波数据样本
        
        Args:
            brain_wave_data: 单个脑波数据对象
            
        Returns:
            处理后的BrainWaveData对象
        """
        try:
            # 使用默认的组合滤波器处理，不指定固定的filter_length
            return self.process_brain_wave_data(brain_wave_data, filter_type='combined')
        except Exception as e:
            # 如果处理失败，返回原始数据
            logger.warning(f"单样本MNE处理失败: {str(e)}")
            return brain_wave_data
    
    def get_filter_info(self) -> Dict:
        """
        获取当前滤波器配置信息
        
        Returns:
            滤波器配置字典
        """
        return {
            'sampling_rate': self.sampling_rate,
            'bandpass_params': self.default_bandpass_params,
            'notch_params': self.default_notch_params,
            'mne_version': mne.__version__
        }
    
    def validate_signal_quality(self, raw: mne.io.Raw) -> Dict:
        """
        验证信号质量
        
        Args:
            raw: MNE Raw对象
            
        Returns:
            信号质量评估结果
        """
        data = raw.get_data()[0]  # 假设单通道
        
        # 计算基本统计指标
        signal_power = np.var(data)
        signal_range = np.ptp(data)  # peak-to-peak
        signal_mean = np.mean(data)
        signal_std = np.std(data)
        
        # 简单的质量评估
        quality_score = 1.0
        issues = []
        
        # 检查信号饱和
        if signal_range > 1000:  # 假设阈值
            quality_score *= 0.8
            issues.append("信号可能饱和")
        
        # 检查直流偏移
        if abs(signal_mean) > 50:  # 假设阈值
            quality_score *= 0.9
            issues.append("存在直流偏移")
        
        # 检查信号变异性
        if signal_std < 1:  # 信号过于平稳
            quality_score *= 0.7
            issues.append("信号变异性过低")
        
        return {
            'quality_score': quality_score,
            'signal_power': signal_power,
            'signal_range': signal_range,
            'signal_mean': signal_mean,
            'signal_std': signal_std,
            'issues': issues
        }